from typing import Dict, List, Optional
from instruction_re.core.datatypes import Instance, Preffix, Relation
from ..instances.Answer import AnswerFormatter
from ..instances.EntityFormatter import EntityFormatter

from ..instances import InstanceFormatter


class RENoICLFormatter(InstanceFormatter):
    """
    Task: Given sentence and entities, type the relation between them
    """

    def format_instance(
        self,
        context: str,
        relations: Optional[List[Dict[str, List[str]]]],
        instruction: str,
        options: List[str],
        examples: List[Relation] = None,
    ) -> Instance:

        instruction = Preffix.INSTRUCTION.value + instruction
        options_joined = ", ".join(options) + "."
        options_joined = Preffix.OPTIONS.value + options_joined
        sentence = Preffix.CONTEXT.value + context
        answer_format = AnswerFormatter.from_relations(relations)
        # answer_format = Preffix.ANSWER.value + answer_format

        entity = EntityFormatter.from_relations(relations)

        relation_triples = [Relation.from_json(r) for r in relations]

        if entity:
            prefix = Preffix.CONDITIONS.value + entity
        else:
            prefix = ""

        instance = Instance(
            instruction=instruction + " " + options_joined,
            context=sentence + " " + prefix + " " + Preffix.ANSWER.value,
            answer=answer_format,
            relations=relation_triples,
        )

        return instance


def main():

    sample = {
        "sentence": "John Wilkes Booth , who assassinated President Lincoln.",
        "relations": [
            {
                "head": {
                    "name": "John Wilkes Booth",
                    "type": "People",
                    "pos": [0, 17],
                },
                "type": "Kill",
                "tail": {
                    "name": "President Lincoln",
                    "type": "People",
                    "pos": [37, 54],
                },
            }
        ],
        "instruction": "Please type the relation between given entities based on the sentence.",
    }

    sample["examples"] = [
        {
            "sentence": "The opera company performed at the Palace of Fine Arts , in San Francisco , on June 30 and July 1-2 , said Kevin O 'Brien , a spokesman for the theater.",
            "relations": [
                {
                    "head": {
                        "name": "Palace of Fine Arts",
                        "type": "Location",
                        "pos": [35, 54],
                    },
                    "type": "Located in",
                    "tail": {
                        "name": "San Francisco",
                        "type": "Location",
                        "pos": [60, 73],
                    },
                }
            ],
        },
    ]

    formatter = RENoICLFormatter()
    print(
        formatter.format_instance(
            context=sample["sentence"],
            relations=sample["relations"],
            examples=sample["examples"],
            instruction=sample["instruction"],
            options=[
                "Organization based in",
                "Located in",
                "Live in",
                "Work for",
                "Kill",
            ],
        )
    )


if __name__ == "__main__":
    main()
